package com.example.xiaosen.config;

import co.elastic.clients.elasticsearch.ElasticsearchClient;
import co.elastic.clients.elasticsearch._types.query_dsl.MatchQuery;
import co.elastic.clients.elasticsearch.core.SearchResponse;
import co.elastic.clients.json.jackson.JacksonJsonpMapper;
import co.elastic.clients.transport.ElasticsearchTransport;
import co.elastic.clients.transport.Version;
import co.elastic.clients.transport.rest_client.RestClientTransport;
import com.fasterxml.jackson.databind.DeserializationFeature;
import com.fasterxml.jackson.databind.ObjectMapper;
import org.elasticsearch.client.RestClient;
import org.springframework.ai.document.Document;
import org.springframework.ai.model.EmbeddingUtils;
import org.springframework.ai.rag.Query;
import org.springframework.ai.rag.retrieval.search.DocumentRetriever;
import org.springframework.ai.rag.retrieval.search.VectorStoreDocumentRetriever;
import org.springframework.ai.vectorstore.SearchRequest;
import org.springframework.ai.vectorstore.VectorStore;
import org.springframework.ai.vectorstore.filter.Filter;
import org.springframework.ai.vectorstore.filter.FilterExpressionTextParser;
import org.springframework.lang.Nullable;
import org.springframework.util.Assert;
import org.springframework.util.StringUtils;

import java.util.List;
import java.util.function.Supplier;
import java.util.stream.Collectors;

/**
 * @Author: zhouxj@ferrotec.com.cn
 * @version: v1.0
 * @Date: 2025/8/14
 * @Description:
 */
public class ESDocumentRetriever implements DocumentRetriever {
    public static final String FILTER_EXPRESSION = "vector_store_filter_expression";
    private final VectorStore vectorStore;
    private final String indexName;
    private final RestClient restClient;
    private final ElasticsearchClient elasticsearchClient;
    private final Double similarityThreshold;
    private final Integer topK;
    private final Supplier<Filter.Expression> filterExpression;

    public ESDocumentRetriever(RestClient restClient, VectorStore vectorStore, String indexName, @Nullable Double similarityThreshold, @Nullable Integer topK,
                               @Nullable Supplier<Filter.Expression> filterExpression) {
        Assert.notNull(vectorStore, "vectorStore cannot be null");
        Assert.notNull(indexName,"indexName must not be null");
        Assert.isTrue(similarityThreshold == null || similarityThreshold >= (double) 0.0F, "similarityThreshold must " +
                "be equal to or greater than 0.0");
        Assert.isTrue(topK == null || topK > 0, "topK must be greater than 0");
        this.vectorStore = vectorStore;
        this.indexName = indexName;
        this.similarityThreshold = similarityThreshold != null ? similarityThreshold : (double) 0.0F;
        this.topK = topK != null ? topK : 4;
        this.filterExpression = filterExpression != null ? filterExpression : () -> null;

        // 初始化ES客户端
        Assert.notNull(restClient, "RestClient must not be null");
        this.restClient = restClient;
        String version = Version.VERSION == null ? "Unknown" : Version.VERSION.toString();
        this.elasticsearchClient = new ElasticsearchClient(
                new RestClientTransport(this.restClient,
                        new JacksonJsonpMapper((new ObjectMapper()).configure(DeserializationFeature.FAIL_ON_UNKNOWN_PROPERTIES, false))))
        .withTransportOptions((t) -> t.addHeader("user-agent", "spring-ai elastic-java/" + version));
    }

    // 实现自定义的混合检索逻辑
    @Override
    public List<Document> retrieve(Query query) {
        Assert.notNull(query, "query cannot be null");
        Filter.Expression requestFilterExpression = this.computeRequestFilterExpression(query);
        // 使用向量数据库查询
        // 这里不直接使用ES的混合检索查询，解耦向量数据库和全文索引库，之后可以切换向量数据库
        SearchRequest searchRequest = SearchRequest.builder()
                .query(query.text())
                .filterExpression(requestFilterExpression)
                .similarityThreshold(this.similarityThreshold)
                .topK(this.topK).build();
        this.vectorStore.similaritySearch(searchRequest);
        // 使用ES进行关键词检索
        // 构建 match 查询
        co.elastic.clients.elasticsearch._types.query_dsl.Query keywordQuery = MatchQuery.of(m -> m
                .field("content") // 替换为你的字段名
                .query(query.text())
        )._toQuery();
        SearchResponse<Document> res = this.elasticsearchClient.search((sr) -> sr.index(this.indexName, , Document.class);
        res.hits().hits().stream().map(this::toDocument).collect(Collectors.toList());
        return this.vectorStore.similaritySearch(searchRequest);
    }

    private Filter.Expression computeRequestFilterExpression(Query query) {
        Object contextFilterExpression = query.context().get("vector_store_filter_expression");
        // 查询上下文中是否包含过滤表达式，否则直接使用文档检索器配置的过滤条件
        return contextFilterExpression != null && StringUtils.hasText(contextFilterExpression.toString()) ?
                (new FilterExpressionTextParser()).parse(contextFilterExpression.toString()) :
                (Filter.Expression) this.filterExpression.get();
    }

    public static Builder builder() {
        return new Builder();
    }

    public static final class Builder {
        private RestClient restClient;
        private VectorStore vectorStore;
        private Double similarityThreshold;
        private String indexName;
        private Integer topK;
        private Supplier<Filter.Expression> filterExpression;

        private Builder() {
        }

        public Builder restClient(RestClient restClient) {
            this.restClient = restClient;
            return this;
        }

        public Builder vectorStore(VectorStore vectorStore) {
            this.vectorStore = vectorStore;
            return this;
        }

        public Builder indexName(String indexName) {
            return this;
        }

        public Builder similarityThreshold(Double similarityThreshold) {
            this.similarityThreshold = similarityThreshold;
            return this;
        }

        public Builder topK(Integer topK) {
            this.topK = topK;
            return this;
        }

        public Builder filterExpression(Filter.Expression filterExpression) {
            this.filterExpression = () -> filterExpression;
            return this;
        }

        public Builder filterExpression(Supplier<Filter.Expression> filterExpression) {
            this.filterExpression = filterExpression;
            return this;
        }

        public ESDocumentRetriever build() {
            return new ESDocumentRetriever(this.restClient, this.vectorStore, this.indexName, this.similarityThreshold, this.topK, this.filterExpression);
        }
    }
}
